package org.apache.lucene.search.payloads;

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.Explanation;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.Searcher;
import org.apache.lucene.search.Similarity;
import org.apache.lucene.search.Weight;
import org.apache.lucene.search.spans.NearSpansOrdered;
import org.apache.lucene.search.spans.NearSpansUnordered;
import org.apache.lucene.search.spans.SpanNearQuery;
import org.apache.lucene.search.spans.SpanQuery;
import org.apache.lucene.search.spans.SpanScorer;
import org.apache.lucene.search.spans.SpanWeight;
import org.apache.lucene.search.spans.Spans;
import org.apache.lucene.util.ToStringUtils;

import java.io.IOException;
import java.util.Collection;
import java.util.Iterator;

/**
 * This class is very similar to
 * {@link org.apache.lucene.search.spans.SpanNearQuery} except that it factors
 * in the value of the payloads located at each of the positions where the
 * {@link org.apache.lucene.search.spans.TermSpans} occurs.
 * <p/>
 * In order to take advantage of this, you must override
 * {@link org.apache.lucene.search.Similarity#scorePayload}
 * which returns 1 by default.
 * <p/>
 * Payload scores are aggregated using a pluggable {@link PayloadFunction}.
 * 
 * @see org.apache.lucene.search.Similarity#scorePayload
 */
public class PayloadNearQuery extends SpanNearQuery {
    protected String fieldName;
    protected PayloadFunction function;

    public PayloadNearQuery(SpanQuery[] clauses, int slop, boolean inOrder) {
        this(clauses, slop, inOrder, new AveragePayloadFunction());
    }

    public PayloadNearQuery(SpanQuery[] clauses, int slop, boolean inOrder, PayloadFunction function) {
        super(clauses, slop, inOrder);
        fieldName = clauses[0].getField(); // all clauses must have same field
        this.function = function;
    }

    @Override
    public Weight createWeight(Searcher searcher) throws IOException {
        return new PayloadNearSpanWeight(this, searcher);
    }

    @Override
    public Object clone() {
        int sz = clauses.size();
        SpanQuery[] newClauses = new SpanQuery[sz];

        for (int i = 0; i < sz; i++) {
            newClauses[i] = (SpanQuery) clauses.get(i).clone();
        }
        PayloadNearQuery boostingNearQuery = new PayloadNearQuery(newClauses, slop, inOrder, function);
        boostingNearQuery.setBoost(getBoost());
        return boostingNearQuery;
    }

    @Override
    public String toString(String field) {
        StringBuilder buffer = new StringBuilder();
        buffer.append("payloadNear([");
        Iterator<SpanQuery> i = clauses.iterator();
        while (i.hasNext()) {
            SpanQuery clause = i.next();
            buffer.append(clause.toString(field));
            if (i.hasNext()) {
                buffer.append(", ");
            }
        }
        buffer.append("], ");
        buffer.append(slop);
        buffer.append(", ");
        buffer.append(inOrder);
        buffer.append(")");
        buffer.append(ToStringUtils.boost(getBoost()));
        return buffer.toString();
    }

    @Override
    public int hashCode() {
        final int prime = 31;
        int result = super.hashCode();
        result = prime * result + ((fieldName == null) ? 0 : fieldName.hashCode());
        result = prime * result + ((function == null) ? 0 : function.hashCode());
        return result;
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj)
            return true;
        if (!super.equals(obj))
            return false;
        if (getClass() != obj.getClass())
            return false;
        PayloadNearQuery other = (PayloadNearQuery) obj;
        if (fieldName == null) {
            if (other.fieldName != null)
                return false;
        } else if (!fieldName.equals(other.fieldName))
            return false;
        if (function == null) {
            if (other.function != null)
                return false;
        } else if (!function.equals(other.function))
            return false;
        return true;
    }

    public class PayloadNearSpanWeight extends SpanWeight {
        public PayloadNearSpanWeight(SpanQuery query, Searcher searcher) throws IOException {
            super(query, searcher);
        }

        @Override
        public Scorer scorer(IndexReader reader, boolean scoreDocsInOrder, boolean topScorer) throws IOException {
            return new PayloadNearSpanScorer(query.getSpans(reader), this, similarity, reader.norms(query.getField()));
        }
    }

    public class PayloadNearSpanScorer extends SpanScorer {
        Spans spans;
        protected float payloadScore;
        private int payloadsSeen;
        Similarity similarity = getSimilarity();

        protected PayloadNearSpanScorer(Spans spans, Weight weight, Similarity similarity, byte[] norms) throws IOException {
            super(spans, weight, similarity, norms);
            this.spans = spans;
        }

        // Get the payloads associated with all underlying subspans
        public void getPayloads(Spans[] subSpans) throws IOException {
            for (int i = 0; i < subSpans.length; i++) {
                if (subSpans[i] instanceof NearSpansOrdered) {
                    if (((NearSpansOrdered) subSpans[i]).isPayloadAvailable()) {
                        processPayloads(((NearSpansOrdered) subSpans[i]).getPayload(), subSpans[i].start(), subSpans[i].end());
                    }
                    getPayloads(((NearSpansOrdered) subSpans[i]).getSubSpans());
                } else if (subSpans[i] instanceof NearSpansUnordered) {
                    if (((NearSpansUnordered) subSpans[i]).isPayloadAvailable()) {
                        processPayloads(((NearSpansUnordered) subSpans[i]).getPayload(), subSpans[i].start(), subSpans[i].end());
                    }
                    getPayloads(((NearSpansUnordered) subSpans[i]).getSubSpans());
                }
            }
        }

        /**
         * By default, uses the {@link PayloadFunction} to score the payloads, but
         * can be overridden to do other things.
         * 
         * @param payLoads The payloads
         * @param start The start position of the span being scored
         * @param end The end position of the span being scored
         * 
         * @see Spans
         */
        protected void processPayloads(Collection<byte[]> payLoads, int start, int end) {
            for (final byte[] thePayload : payLoads) {
                payloadScore = function.currentScore(doc, fieldName, start, end, payloadsSeen, payloadScore, similarity.scorePayload(doc, fieldName, spans.start(), spans.end(), thePayload, 0, thePayload.length));
                ++payloadsSeen;
            }
        }

        //
        @Override
        protected boolean setFreqCurrentDoc() throws IOException {
            if (!more) {
                return false;
            }
            doc = spans.doc();
            freq = 0.0f;
            payloadScore = 0;
            payloadsSeen = 0;
            do {
                int matchLength = spans.end() - spans.start();
                freq += getSimilarity().sloppyFreq(matchLength);
                Spans[] spansArr = new Spans[1];
                spansArr[0] = spans;
                getPayloads(spansArr);
                more = spans.next();
            } while (more && (doc == spans.doc()));
            return true;
        }

        @Override
        public float score() throws IOException {

            return super.score() * function.docScore(doc, fieldName, payloadsSeen, payloadScore);
        }

        @Override
        protected Explanation explain(int doc) throws IOException {
            Explanation result = new Explanation();
            // Add detail about tf/idf...
            Explanation nonPayloadExpl = super.explain(doc);
            result.addDetail(nonPayloadExpl);
            // Add detail about payload
            Explanation payloadExpl = function.explain(doc, payloadsSeen, payloadScore);
            result.addDetail(payloadExpl);
            result.setValue(nonPayloadExpl.getValue() * payloadExpl.getValue());
            result.setDescription("PayloadNearQuery, product of:");
            return result;
        }
    }

}
